import time
import serial
import struct
from enum import Enum, auto
from collections import deque

# 定义常量
HEAD1 = 0xCD
HEAD2 = 0xAB
START_FRAME = 0xABCD
FRAME_SIZE = 16  # 正确帧大小：2(头) + 12(6个int16) + 2(校验)

# 定义状态枚举
class ParserState(Enum):
    WAITING_FOR_HEAD1 = auto()
    WAITING_FOR_HEAD2 = auto()
    WAITING_FOR_PAYLOAD = auto()

# 反馈数据结构
class RikibotFeedback:
    __slots__ = ('start', 'rpmR', 'rpmL', 'batVoltage', 
                 'boardTemp', 'curL_DC', 'curR_DC', 'checksum')
    
    def __init__(self):
        self.start = 0
        self.rpmR = 0
        self.rpmL = 0
        self.batVoltage = 0
        self.boardTemp = 0
        self.curL_DC = 0
        self.curR_DC = 0
        self.checksum = 0

# 命令数据结构
class RikibotCommand:
    __slots__ = ('start', 'mSpeedR', 'mSpeedL', 'checksum')
    
    def __init__(self):
        self.start = START_FRAME
        self.mSpeedR = 0
        self.mSpeedL = 0
        self.checksum = 0

class SerialController:
    def __init__(self):
        # 修改为正确的串口设备
        # ubuntu dev
        #self.port_name = "/dev/ttyUSB0"
        # Windows dev
        self.port_name = "COM6"
        self.baud_rate = 115200
        self.ser = None
        self.parser_state = ParserState.WAITING_FOR_HEAD1
        self.feedback = RikibotFeedback()
        self.command = RikibotCommand()
        self.rx_buffer = bytearray()
        self.frame_queue = deque()
        
    def init_serial(self):
        try:
            self.ser = serial.Serial(
                port=self.port_name,
                baudrate=self.baud_rate,
                bytesize=serial.EIGHTBITS,
                parity=serial.PARITY_NONE,
                stopbits=serial.STOPBITS_ONE,
                timeout=0.01,  # 缩短超时时间
                write_timeout=0.1
            )
            self.ser.reset_input_buffer()
            self.ser.reset_output_buffer()
            print(f"Serial port {self.port_name} opened successfully")
            return True
        except Exception as e:
            print(f"Error opening serial port: {e}")
            return False
    
    def handle_base_data(self):
        """处理并显示基础数据"""
        # 使用ANSI转义码清空当前行并显示新数据
        print(f"\r\033[KRPM: R={self.feedback.rpmR:<5} L={self.feedback.rpmL:<5} | "
              f"Volt: {self.feedback.batVoltage/100:.2f}V | "
              f"Temp: {self.feedback.boardTemp/10:.1f}°C", end='', flush=True)
    
    def parse_frame(self, frame_data):
        """解析完整的数据帧"""
        # 确保帧数据长度正确
        if len(frame_data) != FRAME_SIZE:
            print(f"Invalid frame size: expected {FRAME_SIZE}, got {len(frame_data)}")
            return False
        
        try:
            # 解包数据 - 使用原始字节处理
            # 格式：< 小端字节序
            #       H: 帧起始 (2字节)
            #       6h: 6个short（12字节）
            #       H: 校验和 (2字节)
            unpacked = struct.unpack('<H6hH', frame_data)
            
            # 提取数据
            frame_start = unpacked[0]
            rpmR, rpmL, batVoltage, boardTemp, curL_DC, curR_DC = unpacked[1:7]
            recv_checksum = unpacked[7]
            
            # 验证帧头
            if frame_start != START_FRAME:
                print(f"Invalid frame header: 0x{frame_start:04X} (expected 0x{START_FRAME:04X})")
                return False
            
            # 根据单片机端逻辑计算校验和
            # 直接使用解包后的整数值进行计算
            calc_checksum = (frame_start ^ rpmR ^ rpmL ^ batVoltage ^ 
                             boardTemp ^ curL_DC ^ curR_DC) & 0xFFFF
            
            # 更新反馈数据
            self.feedback.start = frame_start
            self.feedback.rpmR = rpmR
            self.feedback.rpmL = rpmL
            self.feedback.batVoltage = batVoltage
            self.feedback.boardTemp = boardTemp
            self.feedback.curL_DC = curL_DC
            self.feedback.curR_DC = curR_DC
            self.feedback.checksum = recv_checksum
            
            # 验证校验和
            if calc_checksum != recv_checksum:
                print(f"\n\033[KChecksum mismatch: calc 0x{calc_checksum:04X}, recv 0x{recv_checksum:04X}")
                print(f"Details: start=0x{frame_start:04X}, rpmR=0x{rpmR:04X}, rpmL=0x{rpmL:04X}, "
                      f"batVoltage=0x{batVoltage:04X}, boardTemp=0x{boardTemp:04X}, "
                      f"curL_DC=0x{curL_DC:04X}, curR_DC=0x{curR_DC:04X}")
            else:
                # 校验和正确时处理数据
                self.handle_base_data()
            
            return True
        except struct.error as e:
            print(f"Struct unpack error: {e}")
            return False
    
    def process_rx_data(self, new_data):
        """处理接收到的数据"""
        if not new_data:
            return
            
        self.rx_buffer.extend(new_data)
        
        while len(self.rx_buffer) >= FRAME_SIZE:
            # 查找帧头
            found_frame = False
            for i in range(len(self.rx_buffer) - FRAME_SIZE + 1):
                if self.rx_buffer[i] == HEAD1 and self.rx_buffer[i+1] == HEAD2:
                    # 找到可能的帧头
                    frame_data = bytes(self.rx_buffer[i:i+FRAME_SIZE])
                    if self.parse_frame(frame_data):
                        found_frame = True
                        # 移除处理过的数据
                        self.rx_buffer = self.rx_buffer[i+FRAME_SIZE:]
                        break
                    else:
                        # 解析失败，继续搜索下一个帧头
                        continue
            
            # 如果没有找到完整帧，清空缓冲区
            if not found_frame:
                # 打印前20个字节帮助调试
                if len(self.rx_buffer) > 0:
                    print(f"\nUnprocessed data: {bytes(self.rx_buffer[:min(20, len(self.rx_buffer))]).hex()}")
                self.rx_buffer.clear()
                break
    
    def read_serial_data(self):
        """读取串口数据"""
        try:
            # 批量读取可用数据
            if self.ser.in_waiting > 0:
                new_data = self.ser.read(self.ser.in_waiting)
                self.process_rx_data(new_data)
        except serial.SerialException as e:
            print(f"\nSerial read error: {e}")
            # 尝试重新初始化串口
            try:
                if self.ser.is_open:
                    self.ser.close()
                time.sleep(0.5)
                self.ser.open()
                self.ser.reset_input_buffer()
                print("Serial port reinitialized")
            except Exception as e2:
                print(f"Reinitialization failed: {e2}")
    
    def send_speed(self, left_speed, right_speed):
        """发送速度命令"""
        # 确保速度值在有效范围内
        self.command.mSpeedR = max(min(right_speed, 32767), -32768)
        self.command.mSpeedL = max(min(left_speed, 32767), -32768)
        
        # 计算校验和并确保为无符号值
        checksum = self.command.start ^ self.command.mSpeedR ^ self.command.mSpeedL
        self.command.checksum = checksum & 0xFFFF  # 确保在0-65535范围内
        
        # 打包数据为二进制格式
        packed_data = struct.pack('<HhhH',
                                  self.command.start,
                                  self.command.mSpeedR,
                                  self.command.mSpeedL,
                                  self.command.checksum)
        try:
            self.ser.write(packed_data)
        except Exception as e:
            print(f"Send error: {e}")

def main():
    controller = SerialController()
    if not controller.init_serial():
        print("Serial port initialization failed")
        return
    
    print("Starting Rikibot controller. Press Ctrl+C to exit...")
    
    try:
        while True:
            start_time = time.time()
            
            # 读取并处理串口数据
            controller.read_serial_data()
            
            # 发送速度命令 (20ms周期)
            controller.send_speed(100, -100)
            
            # 精确控制循环时间
            elapsed = time.time() - start_time
            sleep_time = max(0.0, 0.02 - elapsed)
            time.sleep(sleep_time)
    
    except KeyboardInterrupt:
        print("\nProgram terminated by user")
    finally:
        if controller.ser and controller.ser.is_open:
            controller.ser.close()
            print("\nSerial port closed")

if __name__ == "__main__":
    main()